import numpy as np
import time
import matplotlib.pyplot as plt
#定义函数f(x)
def f(x):
    return np.sin(100*x+np.pi/3)+0.001*x**2+x-5
#离散傅里叶变换,求Ck
def DFT(y):
    N=len(y)
    c=np.zeros(N,dtype=np.complex128)#复数类型
    i=np.array([a for a in range(N)])#i属于[0,N-1]
    for k in range(N):
        #求和公式的各个指数化成一个列表
        e=np.exp(-2j*np.pi*i*k/N)
        e.reshape((N*1))
        #计算内积
        c[k]=np.dot(y,e)
    
    return c
#定义插值函数s(x)【即傅里叶逆变换】，取不同的x=i*2Π/N对应fi
def s(x,c):
    N=len(c)
    #k属于[0,N-1]
    k=np.array([a for a in range(N)])
    s_=np.zeros(N,dtype=np.complex128)
    #与DFT类似
    for i,x_ in enumerate(x):
        e=np.exp(1j*k*x_)
        e.reshape((N,1))
        s_[i]=np.dot(c,e)
    
    return s_
#快速傅里叶变换    
def FFT(y):
    N=len(y)
    #当y只有一个数时直接返回
    if N <= 1:
        return y
    #偶数下标的列表，从0开始按两位选取
    even=FFT(y[0::2])
    #奇数下标的列表，从1开始按两位选取
    odd=FFT(y[1::2])
    #求奇数下标部分的傅里叶变换
    T=[np.exp(-2j*np.pi*k/N) * odd[k] for k in range(N//2)]
    #通过蝶形运算，将<N/2的偶数部分加上T，将>N/2的偶数部分（根据性质变量可化为<N/2）减去T，即可求得DFT值
    return np.array( [even[k]+T[k] for k in range(N//2)] + [even[k]-T[k] for k in range(N//2)] )    
#初始化数据
N=16
#定义x为序列{2Π/N*i}
x=np.array([2*np.pi/N*i for i in range(N)])
x_=np.linspace(0,2*np.pi,50)
w=np.array([a for a in range(N)])
y=f(x)
y_=f(x_)
count=0
#起始时间
time1=time.time()
dft_result=DFT(y)
#得到的是X(ejw)
print("DFT:\n",dft_result)
#传入的是傅里叶系数，是X(ejw)/N
S=s(x,dft_result/N)
while count<=10000000:
    count+=1
#DFT结束时间
time2=time.time()
fft_result=FFT(y)
print("FFT:\n",fft_result)
S2=s(x,fft_result/N)
#两个步骤消耗一样时间，方便比较
while count>=0:
    count-=1
#FFT结束时间
time3=time.time()
#画图
# 创建画布和子图
fig,ax=plt.subplots(3,2,figsize=(10,6))
ax[0][0].plot(x_,y_,c='r')
ax[0][0].set_title('Init:f(x)')
ax[0][1].scatter(x,y)
ax[0][1].set_title('fn')
ax[1][0].plot(x,S,c='b')
ax[1][0].set_title('DFT:S(x)')
ax[1][1].stem(w,dft_result)
ax[1][1].set_title('DFT:X(jw)')
ax[2][0].plot(x,S2,c='g')
ax[2][0].set_title('FFT:S\'(x)')
ax[2][1].stem(w,fft_result)
ax[2][1].set_title('FFT:X(jw)')
# 调整子图间的间距
plt.tight_layout()
plt.show()
print("DFT Used time:",time2-time1,"s")
print("FFT Used time:",time3-time2,"s")